"""
MomentumStrategy_Optimized.py
ETF动量轮动策略

该策略通过计算多个ETF的历史动量，选择表现最好的ETF进行持有，
定期轮动以获得超额收益。

主要逻辑：
1. 计算每个ETF的N日动量
2. 选择动量最强的ETF进行持有
3. 定期轮动，扣除交易成本
4. 当所有ETF动量都为负时，空仓观望
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from function import *

# ==================== 策略配置参数 ====================
class StrategyConfig:
    """策略配置参数"""
    START_TIME = '20190101'  # 策略开始时间
    MOMENTUM_DAYS = 10  # 动量计算周期（交易日）
    TRADE_RATE = 0.86 / 10000  # 交易费率（万分之0.86）
    TOP_N_HOLDINGS = 1  # 持仓动量前N个ETF

    # ==================== 调仓周期配置 ====================
    REBALANCE_FREQUENCY = 10  # 调仓频率（交易日），0表示每天调仓
    REBALANCE_MODE = 'fixed'  # 调仓模式：'fixed'(定期), 'dynamic'(动态), 'hybrid'(混合)
    MOMENTUM_THRESHOLD = 0.05  # 动量变化阈值（动态调仓时使用）

    # 参与轮动的ETF列表
    ETF_FILES = [
        '159915.csv',  # 创业板ETF
        '159939.csv',  # 信息技术
        '159941.csv',  # 纳斯达克
        '159755.csv',  # 电池
        #'512100.csv',  # 中证1000
        '512690.csv',  # 白酒
        '512890.csv',  # 红利低波
        #'515170.csv',  # 食品饮料
        #'510300.csv',  # 华泰柏瑞沪深300ETF
        #'510880.csv',  # 红利
        '518880.csv',  # 华安黄金ETF
        '159995.csv', #芯片
        '562500.csv', #机器人
        #'513180.csv',#恒生科技,
        '159992.csv', #创新药
        '515220.csv',
    ]

    # 数据文件路径
    DATA_PATH = 'data/'

    # 输出文件路径
    OUTPUT_FILE = 'rawData/MomentumStrategy_Optimized_result.csv'

    # 绘图配置
    PLOT_TOP_N = 5  # 只显示前N个ETF的净值曲线

# ==================== 数据处理函数 ====================
def load_etf_data(etf_files, data_path):
    """
    加载ETF数据并进行预处理

    Args:
        etf_files (list): ETF文件名列表
        data_path (str): 数据文件路径

    Returns:
        pd.DataFrame: 合并后的ETF数据
    """
    # 以第一个ETF为基准时间轴
    first_etf_code = etf_files[0].split('.')[0]
    df = pd.read_csv(f'{data_path}{etf_files[0]}', encoding='utf-8', parse_dates=['date'])
    df.rename(columns={'date': 'candle_end_time'}, inplace=True)

    # 对基准ETF也应用时间过滤
    df = df[df['candle_end_time'] >= pd.to_datetime(StrategyConfig.START_TIME)]

    # 合并所有ETF数据
    for etf_file in etf_files:
        etf_code = etf_file.split('.')[0]
        etf_data = load_single_etf(etf_file, data_path, etf_code)
        df = pd.merge(df, etf_data, how='outer', on=['candle_end_time'])

    # 按时间排序，确保时间序列正确
    df = df.sort_values('candle_end_time').reset_index(drop=True)

    return df

def load_single_etf(etf_file, data_path, etf_code):
    """
    加载单个ETF数据并进行计算

    Args:
        etf_file (str): ETF文件名
        data_path (str): 数据文件路径
        etf_code (str): ETF代码

    Returns:
        pd.DataFrame: 处理后的ETF数据
    """
    # 读取原始数据
    df = pd.read_csv(f'{data_path}{etf_file}', encoding='utf-8', parse_dates=['date'])

    # 只过滤策略开始时间之后的数据（如果ETF有更早的数据，会被保留）
    df = df[df['date'] >= pd.to_datetime(StrategyConfig.START_TIME)]

    # 生成开盘价（用前一天收盘价）
    df['open'] = df['close_price'].shift(1)
    df['close'] = df['close_price']

    # 计算日收益率
    df[f'{etf_code}_amp'] = df['close'] / df['close'].shift(1) - 1

    # 重命名列
    df.rename(columns={
        'date': 'candle_end_time',
        'open': f'{etf_code}_open',
        'close': f'{etf_code}_close'
    }, inplace=True)

    # 选择需要的列
    return df[['candle_end_time', f'{etf_code}_open', f'{etf_code}_close', f'{etf_code}_amp']]

def calculate_momentum(df, etf_files):
    """
    计算各ETF的动量指标

    Args:
        df (pd.DataFrame): ETF数据
        etf_files (list): ETF文件名列表

    Returns:
        pd.DataFrame: 包含动量指标的数据
    """
    for etf_file in etf_files:
        etf_code = etf_file.split('.')[0]
        # 计算N日动量（收益率）
        df[f'{etf_code}_mom'] = df[f'{etf_code}_close'].pct_change(periods=StrategyConfig.MOMENTUM_DAYS)
    return df

# ==================== 策略逻辑函数 ====================
def generate_signals(df, etf_files):
    """
    生成交易信号：选择动量前N个ETF，支持定期调仓

    Args:
        df (pd.DataFrame): ETF数据
        etf_files (list): ETF文件名列表

    Returns:
        pd.DataFrame: 包含交易信号的数据
    """
    # 获取所有ETF代码
    etf_codes = [file.split('.')[0] for file in etf_files]
    momentum_cols = [f'{code}_mom' for code in etf_codes]

    # 为每一行生成信号：选择动量前N个ETF
    for idx in df.index:
        # 获取当日的动量数据
        row_momentum = df.loc[idx, momentum_cols]

        # 按动量降序排序
        sorted_momentum = row_momentum.sort_values(ascending=False)

        # 选择前N个正动量的ETF
        positive_momentum = sorted_momentum[sorted_momentum > 0]

        if len(positive_momentum) == 0:
            # 所有ETF动量都为负，空仓
            df.loc[idx, 'signal'] = 'empty'
        else:
            # 选择前N个正动量的ETF
            top_n = min(StrategyConfig.TOP_N_HOLDINGS, len(positive_momentum))
            selected_etfs = positive_momentum.head(top_n).index.tolist()

            # 将ETF代码从列名转换回代码
            selected_codes = [col.replace('_mom', '') for col in selected_etfs]

            # 用逗号分隔多个ETF
            df.loc[idx, 'signal'] = ','.join(selected_codes)

    # 根据调仓配置处理信号
    if StrategyConfig.REBALANCE_FREQUENCY == 0:
        # 每天调仓（原逻辑）
        df['signal'] = df['signal'].ffill()
    else:
        # 定期调仓
        df = apply_rebalancing_schedule(df, etf_files)

    # 实际持仓比信号晚一天（收盘后确定信号，次日执行）
    df['position'] = df['signal'].shift(1)

    # 删除持仓为空的数据
    df.dropna(subset=['position'], inplace=True)
    return df

def apply_rebalancing_schedule(df, etf_files):
    """
    根据调仓周期应用调仓计划

    Args:
        df (pd.DataFrame): 包含信号的数据
        etf_files (list): ETF文件名列表

    Returns:
        pd.DataFrame: 处理后的数据
    """
    if StrategyConfig.REBALANCE_MODE == 'fixed':
        return apply_fixed_rebalancing(df)
    elif StrategyConfig.REBALANCE_MODE == 'dynamic':
        return apply_dynamic_rebalancing(df, etf_files)
    elif StrategyConfig.REBALANCE_MODE == 'hybrid':
        return apply_hybrid_rebalancing(df, etf_files)
    else:
        raise ValueError(f"未知的调仓模式: {StrategyConfig.REBALANCE_MODE}")

def apply_fixed_rebalancing(df):
    """
    应用定期调仓策略

    Args:
        df (pd.DataFrame): 包含信号的数据

    Returns:
        pd.DataFrame: 处理后的数据
    """
    # 确定调仓日期
    df['day_index'] = range(len(df))
    df['rebalance_day'] = (df['day_index'] % StrategyConfig.REBALANCE_FREQUENCY == 0)

    # 只在调仓日更新信号
    df['signal_restricted'] = df['signal'].where(df['rebalance_day'])

    # 前向填充信号，保持持仓不变
    df['signal'] = df['signal_restricted'].ffill()

    # 清理临时列
    df.drop(['day_index', 'rebalance_day', 'signal_restricted'], axis=1, inplace=True)

    return df

def apply_dynamic_rebalancing(df, etf_files):
    """
    应用动态调仓策略（基于动量变化阈值）

    Args:
        df (pd.DataFrame): 包含信号的数据
        etf_files (list): ETF文件名列表

    Returns:
        pd.DataFrame: 处理后的数据
    """
    # 计算持仓ETF的平均动量变化
    etf_codes = [file.split('.')[0] for file in etf_files]

    df['should_rebalance'] = False

    for i in range(1, len(df)):
        current_signal = df.loc[i, 'signal']
        previous_signal = df.loc[i-1, 'signal']

        if current_signal != previous_signal:
            if current_signal == 'empty' or previous_signal == 'empty':
                # 空仓和有仓位之间的变化总是调仓
                df.loc[i, 'should_rebalance'] = True
            else:
                # 计算动量变化幅度
                current_etfs = current_signal.split(',')
                previous_etfs = previous_signal.split(',')

                # 如果持仓ETF发生变化，计算动量差异
                if set(current_etfs) != set(previous_etfs):
                    momentum_change = calculate_momentum_change(
                        df, i, current_etfs, previous_etfs, etf_codes
                    )
                    if momentum_change > StrategyConfig.MOMENTUM_THRESHOLD:
                        df.loc[i, 'should_rebalance'] = True
                else:
                    # 持仓ETF相同，检查动量是否显著变化
                    momentum_change = calculate_single_etf_momentum_change(
                        df, i, current_etfs[0], etf_codes
                    )
                    if momentum_change > StrategyConfig.MOMENTUM_THRESHOLD:
                        df.loc[i, 'should_rebalance'] = True

    # 添加定期强制调仓（确保不会长期不调仓）
    df['day_index'] = range(len(df))
    df['forced_rebalance'] = (df['day_index'] % StrategyConfig.REBALANCE_FREQUENCY == 0)
    df['should_rebalance'] = df['should_rebalance'] | df['forced_rebalance']

    # 应用调仓信号
    df['signal_restricted'] = df['signal'].where(df['should_rebalance'])
    df['signal'] = df['signal_restricted'].ffill()

    # 清理临时列
    df.drop(['day_index', 'forced_rebalance', 'should_rebalance', 'signal_restricted'], axis=1, inplace=True)

    return df

def apply_hybrid_rebalancing(df, etf_files):
    """
    应用混合调仓策略（定期+动态）

    Args:
        df (pd.DataFrame): 包含信号的数据
        etf_files (list): ETF文件名列表

    Returns:
        pd.DataFrame: 处理后的数据
    """
    # 先应用动态调仓
    df = apply_dynamic_rebalancing(df, etf_files)
    return df

def calculate_momentum_change(df, idx, current_etfs, previous_etfs, etf_codes):
    """
    计算持仓ETF的动量变化幅度

    Args:
        df (pd.DataFrame): 数据
        idx (int): 当前索引
        current_etfs (list): 当前持仓ETF
        previous_etfs (list): 之前持仓ETF
        etf_codes (list): 所有ETF代码

    Returns:
        float: 动量变化幅度
    """
    current_momentum = 0
    previous_momentum = 0

    # 计算当前持仓ETF的平均动量
    for etf in current_etfs:
        if etf in etf_codes and f'{etf}_mom' in df.columns:
            current_momentum += df.loc[idx, f'{etf}_mom']
    current_momentum /= len(current_etfs) if current_etfs else 1

    # 计算之前持仓ETF的平均动量
    for etf in previous_etfs:
        if etf in etf_codes and f'{etf}_mom' in df.columns:
            previous_momentum += df.loc[idx, f'{etf}_mom']
    previous_momentum /= len(previous_etfs) if previous_etfs else 1

    return abs(current_momentum - previous_momentum)

def calculate_single_etf_momentum_change(df, idx, etf_code, etf_codes):
    """
    计算单个ETF的动量变化幅度

    Args:
        df (pd.DataFrame): 数据
        idx (int): 当前索引
        etf_code (str): ETF代码
        etf_codes (list): 所有ETF代码

    Returns:
        float: 动量变化幅度
    """
    if etf_code not in etf_codes or f'{etf_code}_mom' not in df.columns:
        return 0

    # 计算该ETF相对于其他ETF的动量排名变化
    current_momentum = df.loc[idx, f'{etf_code}_mom']

    if idx > 0:
        previous_momentum = df.loc[idx-1, f'{etf_code}_mom']
        return abs(current_momentum - previous_momentum)

    return 0

def calculate_strategy_returns(df, etf_files):
    """
    计算策略收益率（考虑交易成本，支持多ETF平均分配仓位）

    Args:
        df (pd.DataFrame): ETF数据
        etf_files (list): ETF文件名列表

    Returns:
        pd.DataFrame: 包含策略收益率的数据
    """
    # 计算策略每日收益率（不考虑交易成本）
    df['strategy_amp'] = 0.0

    for idx in df.index:
        position = df.loc[idx, 'position']

        if position == 'empty':
            # 空仓
            df.loc[idx, 'strategy_amp'] = 0
        else:
            # 多ETF持仓，平均分配仓位
            held_etfs = position.split(',')
            weight = 1.0 / len(held_etfs)  # 平均权重

            daily_return = 0.0
            for etf_code in held_etfs:
                if etf_code in [file.split('.')[0] for file in etf_files]:
                    daily_return += weight * df.loc[idx, f'{etf_code}_amp']

            df.loc[idx, 'strategy_amp'] = daily_return

    # 标记调仓日
    df.loc[df['position'] != df['position'].shift(1), 'trade_time'] = df['candle_end_time']

    # 调整调仓日的收益率（考虑交易成本）
    df['strategy_amp_adjust'] = df['strategy_amp']

    for idx in df[df['trade_time'].notnull()].index:
        position = df.loc[idx, 'position']

        if position == 'empty':
            # 空仓
            df.loc[idx, 'strategy_amp_adjust'] = 0
        else:
            # 计算买入成本
            held_etfs = position.split(',')
            weight = 1.0 / len(held_etfs)

            adjusted_return = 0.0
            for etf_code in held_etfs:
                if etf_code in [file.split('.')[0] for file in etf_files]:
                    # 考虑买入成本
                    cost_adjusted_return = (
                        df.loc[idx, f'{etf_code}_close'] /
                        (df.loc[idx, f'{etf_code}_open'] * (1 + StrategyConfig.TRADE_RATE)) - 1
                    )
                    adjusted_return += weight * cost_adjusted_return

            df.loc[idx, 'strategy_amp_adjust'] = adjusted_return

    # 卖出成本（次日有调仓时）
    sell_condition = (df['trade_time'].shift(-1).notnull()) & (df['position'] != 'empty')
    df.loc[sell_condition, 'strategy_amp_adjust'] = (
        (1 + df['strategy_amp']) * (1 - StrategyConfig.TRADE_RATE) - 1
    )

    # 清理中间变量
    df.drop(['strategy_amp', 'signal'], axis=1, inplace=True)
    df.reset_index(drop=True, inplace=True)
    return df

def analyze_rebalancing_statistics(df, etf_files):
    """
    分析调仓统计信息

    Args:
        df (pd.DataFrame): 策略结果数据
        etf_files (list): ETF文件名列表

    Returns:
        dict: 调仓统计信息
    """
    if 'trade_time' not in df.columns:
        return {"error": "未找到调仓时间数据"}

    # 获取所有调仓日期
    rebalance_dates = df[df['trade_time'].notnull()]['candle_end_time'].tolist()

    # 计算调仓间隔
    if len(rebalance_dates) > 1:
        rebalance_intervals = []
        for i in range(1, len(rebalance_dates)):
            interval = (rebalance_dates[i] - rebalance_dates[i-1]).days
            rebalance_intervals.append(interval)

        avg_interval = sum(rebalance_intervals) / len(rebalance_intervals)
        min_interval = min(rebalance_intervals)
        max_interval = max(rebalance_intervals)
    else:
        avg_interval = min_interval = max_interval = 0

    # 统计调仓次数
    total_rebalances = len(rebalance_dates)

    # 计算平均持仓时间
    if total_rebalances > 0:
        total_days = (df['candle_end_time'].max() - df['candle_end_time'].min()).days
        avg_holding_days = total_days / total_rebalances if total_rebalances > 0 else 0
    else:
        avg_holding_days = 0

    # 统计各ETF的持仓次数
    etf_holdings = {}
    etf_codes = [file.split('.')[0] for file in etf_files]

    for etf_code in etf_codes:
        holdings_count = len(df[df['position'].str.contains(etf_code, na=False)])
        etf_holdings[etf_code] = holdings_count

    # 计算空仓时间占比
    empty_days = len(df[df['position'] == 'empty'])
    empty_ratio = empty_days / len(df) if len(df) > 0 else 0

    # 计算交易成本
    if total_rebalances > 0:
        total_cost = total_rebalances * 2 * StrategyConfig.TRADE_RATE  # 买入+卖出
        cost_ratio = total_cost * 100  # 转换为百分比
    else:
        total_cost = cost_ratio = 0

    return {
        "总调仓次数": total_rebalances,
        "平均调仓间隔天数": f"{avg_interval:.1f}",
        "最短调仓间隔": f"{min_interval}",
        "最长调仓间隔": f"{max_interval}",
        "平均持仓天数": f"{avg_holding_days:.1f}",
        "空仓时间占比": f"{empty_ratio*100:.1f}%",
        "累计交易成本": f"{total_cost*10000:.2f}bp",
        "各ETF持仓次数": etf_holdings,
        "调仓日期列表": [str(date) for date in rebalance_dates[:10]]  # 只显示前10个
    }

def calculate_net_values(df, etf_files):
    """
    计算净值曲线

    Args:
        df (pd.DataFrame): ETF数据
        etf_files (list): ETF文件名列表

    Returns:
        pd.DataFrame: 包含净值曲线的数据
    """
    # 计算各ETF的净值
    for etf_file in etf_files:
        etf_code = etf_file.split('.')[0]
        # 获取第一个有效价格（跳过NaN值）
        first_valid_idx = df[f'{etf_code}_close'].first_valid_index()
        if first_valid_idx is not None:
            first_close = df[f'{etf_code}_close'].loc[first_valid_idx]
            df[f'{etf_code}_net'] = df[f'{etf_code}_close'] / first_close

    # 计算策略净值
    df['strategy_net'] = (1 + df['strategy_amp_adjust']).cumprod()
    return df

# ==================== 可视化函数 ====================
def plot_performance(df, etf_files):
    """
    绘制策略表现图表

    Args:
        df (pd.DataFrame): ETF数据
        etf_files (list): ETF文件名列表
    """
    plt.figure(figsize=(12, 8))

    # 绘制策略净值曲线
    plt.plot(df['candle_end_time'], df['strategy_net'],
             label='momentum strategy', linewidth=2, color='red')

    # 绘制所有ETF的净值曲线
    for etf_file in etf_files:
        etf_code = etf_file.split('.')[0]

        # 检查数据列是否存在且有有效数据
        if f'{etf_code}_net' in df.columns and not df[f'{etf_code}_net'].isna().all():
            plt.plot(df['candle_end_time'], df[f'{etf_code}_net'],
                    label=f'{etf_code}_净值', alpha=0.7)

    plt.title('momentum strategy', fontsize=14)
    plt.xlabel('date', fontsize=12)
    plt.ylabel('net', fontsize=12)
    plt.legend(fontsize=10)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()

    # 保存图片到image目录
    plt.savefig('image/etf_momentum_strategy.png', dpi=300, bbox_inches='tight')
    plt.show()

# ==================== 主函数 ====================
def main(etf_files):
    """
    主函数：执行ETF动量轮动策略

    Args:
        etf_files (list): ETF文件名列表
    """
    print("开始执行ETF动量轮动策略...")
    print(f"参与轮动的ETF数量: {len(etf_files)}")
    print(f"动量计算周期: {StrategyConfig.MOMENTUM_DAYS}个交易日")
    print(f"交易费率: {StrategyConfig.TRADE_RATE * 10000:.2f}bp")
    print("-" * 50)

    # 1. 加载数据
    print("正在加载ETF数据...")
    df = load_etf_data(etf_files, StrategyConfig.DATA_PATH)

    # 2. 计算动量
    print("正在计算动量指标...")
    df = calculate_momentum(df, etf_files)

    # 3. 生成交易信号
    print("正在生成交易信号...")
    df = generate_signals(df, etf_files)

    # 4. 计算策略收益率
    print("正在计算策略收益率...")
    df = calculate_strategy_returns(df, etf_files)

    # 5. 计算净值
    print("正在计算净值曲线...")
    df = calculate_net_values(df, etf_files)

    # 6. 评估策略
    print("正在评估策略表现...")
    evaluation_results = evaluate_investment(df, 'strategy_net', time='candle_end_time')
    print("\n策略评估结果:")
    print(evaluation_results)

    # 显示调仓统计信息
    print("\n正在分析调仓统计信息...")
    rebalance_stats = analyze_rebalancing_statistics(df, etf_files)
    print("\n调仓统计信息:")
    for key, value in rebalance_stats.items():
        if key != "各ETF持仓次数" and key != "调仓日期列表":
            print(f"  {key}: {value}")

    # 显示ETF持仓统计（前5个）
    print("\n  各ETF持仓次数（前5个）:")
    etf_holdings = rebalance_stats.get("各ETF持仓次数", {})
    sorted_holdings = sorted(etf_holdings.items(), key=lambda x: x[1], reverse=True)
    for etf_code, count in sorted_holdings[:5]:
        print(f"    {etf_code}: {count}天")

    # 显示年度收益详情
    yearly_returns_str = evaluation_results.loc['年度收益详情', 0]
    if yearly_returns_str != '{}' and yearly_returns_str != 'None':
        yearly_returns = eval(yearly_returns_str)
        print("\n年度收益率详情:")
        for year, return_rate in yearly_returns.items():
            print(f"  {year}年: {return_rate:+.2f}%")

        # 计算统计信息
        returns_list = list(yearly_returns.values())
        positive_years = sum(1 for r in returns_list if r > 0)
        negative_years = sum(1 for r in returns_list if r < 0)
        avg_return = sum(returns_list) / len(returns_list) if returns_list else 0
        max_return = max(returns_list) if returns_list else 0
        min_return = min(returns_list) if returns_list else 0

        print(f"\n年度收益统计:")
        print(f"  盈利年数: {positive_years}/{len(returns_list)}")
        print(f"  亏损年数: {negative_years}/{len(returns_list)}")
        print(f"  平均年收益: {avg_return:+.2f}%")
        print(f"  最高年收益: {max_return:+.2f}%")
        print(f"  最低年收益: {min_return:+.2f}%")

    # 显示配置信息
    print(f"\n当前策略配置:")
    print(f"  调仓频率: {StrategyConfig.REBALANCE_FREQUENCY}交易日")
    print(f"  调仓模式: {StrategyConfig.REBALANCE_MODE}")
    if StrategyConfig.REBALANCE_MODE in ['dynamic', 'hybrid']:
        print(f"  动量阈值: {StrategyConfig.MOMENTUM_THRESHOLD}")
    print(f"  动量计算周期: {StrategyConfig.MOMENTUM_DAYS}交易日")
    print(f"  持仓ETF数量: {StrategyConfig.TOP_N_HOLDINGS}个")

    # 7. 绘制图表
    print("正在绘制表现图表...")
    plot_performance(df, etf_files)

    # 8. 保存结果
    print(f"正在保存结果到 {StrategyConfig.OUTPUT_FILE}...")

    # 创建输出结果DataFrame
    etf_codes = [file.split('.')[0] for file in etf_files]
    output_df = pd.DataFrame()

    # 添加日期
    output_df['candle_end_time'] = df['candle_end_time']

    # 添加所有ETF的动量收益率
    for etf_code in etf_codes:
        output_df[f'{etf_code}_动量收益率'] = df[f'{etf_code}_mom']

    # 添加持仓ETF
    output_df['持仓ETF'] = df['position']

    # 添加持仓数量
    output_df['持仓数量'] = df['position'].apply(lambda x: 0 if x == 'empty' else len(x.split(',')))

    # 添加当前净值
    output_df['当前净值'] = df['strategy_net']

    # 添加持仓收益率
    output_df['持仓收益率'] = df['strategy_amp_adjust']

    # 添加调仓标记
    output_df['是否调仓'] = df['trade_time'].notna().astype(int)

    # 为每个持仓ETF添加单独的列，显示是否持仓
    for etf_code in etf_codes:
        output_df[f'持仓_{etf_code}'] = df['position'].apply(
            lambda x: 1 if etf_code in x.split(',') else 0
        )

    # 保存增强的结果
    output_df.to_csv(StrategyConfig.OUTPUT_FILE, encoding='utf-8', index=False)

    print("\n策略执行完成！")

if __name__ == '__main__':
    # 设置pandas显示选项
    pd.set_option('expand_frame_repr', False)
    pd.set_option('display.max_rows', 5000)

    # 执行策略
    main(StrategyConfig.ETF_FILES)